import moderngl
import numpy as np
import pygame
from pygame.locals import OPENGL, DOUBLEBUF

pygame.init()

vertex_shader = '''
#version 430
in vec2 position;
in float offset;

out vec4 color;

void main() {
    color = vec4(offset, offset, offset, 1.0);
    gl_Position = vec4(position, 0.0, 1.0);
}
'''

fragment_shader = '''
#version 430
in vec4 color;

void main() {
    gl_FragColor = color;
}
'''

compute_shader = '''
#version 430
layout(local_size_x=8, local_size_y=8) in;

layout(std430, binding=0) buffer Offset
{
    float offsets[];
};

void main()
{
    vec2 pos = gl_GlobalInvocationID.xy - vec2(256, 256);
    float r = length(pos);
    uint i = gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;
    if (r <= 255) {
        offsets[i] = r / 255.0;
    } else {
        offsets[i] = 1.0;
    }
}
'''


class Window:
    def __init__(self):
        size = 512
        half_size = size / 2
        self.width = size
        self.height = size
        pygame.display.set_mode((self.width, self.height), DOUBLEBUF | OPENGL)
        self.ctx = moderngl.create_context()
        self.program = self.ctx.program(
            vertex_shader=vertex_shader,
            fragment_shader=fragment_shader,
        )
        data = np.ones(size * size, dtype='f4')
        self.buffer = self.ctx.buffer(data)
        positions = []
        for x in range(size):
            for y in range(size):
                positions.append((x - half_size, y - half_size))
        position: np.ndarray = np.array(positions, dtype='f4') / half_size
        vbo = self.ctx.buffer(position.flatten())
        self.vao = self.ctx.vertex_array(self.program, [
            (vbo, '2f', 'position'),
            (self.buffer, 'f', 'offset'),
        ])

        self.compute_shader = self.ctx.compute_shader(compute_shader)

    def update(self, dt: float, wt: float, events):
        self.buffer.bind_to_storage_buffer(0)
        self.compute_shader.run(group_x=64, group_y=64)

    def render(self):
        self.vao.render(mode=moderngl.POINTS)
        pygame.display.flip()

    @staticmethod
    def quit():
        pygame.quit()

    def run(self, fps: int = 60):
        world_time: int = 0
        clock = pygame.time.Clock()
        while True:
            events = pygame.event.get()
            for event in events:
                if event.type == pygame.QUIT:
                    return self.quit()
                if event.type == pygame.KEYDOWN:
                    if event.key == 27:
                        return self.quit()
            dt = clock.get_time() / 1000
            world_time += dt
            self.update(dt, world_time, events)
            self.render()
            print(clock.get_fps())
            clock.tick(fps)

# 给一个数组,快速排序
    def quickSort(arr):
        if len(arr) <= 1:
            return arr
        else:
            pivot = arr[0]
            less = [i for i in arr[1:] if i <= pivot]
            greater = [i for i in arr[1:] if i > pivot]
            return quickSort(less) + [pivot] + quickSort(greater)


if __name__ == '__main__':
    Window().run(fps=60)
